import requests
import os
import time
from typing import Optional, Dict, Any
from src.processors.platform_apis import get_platform_config


class TikHubSora2VideoGenerator:
    """
    TikHub Sora2视频生成器
    用于处理图片上传、视频生成任务创建、状态查询和获取无水印视频详情
    """

    def __init__(self, api_key: Optional[str] = None):
        """
        初始化TikHub Sora2视频生成器
        
        Args:
            api_key (str, optional): TikHub API密钥
        """
        self.platform = "tikhub"
        self.platform_config = get_platform_config(self.platform)
        self.api_key = api_key or os.environ.get("TIKHUB_API_KEY")
        self.base_url = self.platform_config.base_url if self.platform_config else "https://api.tikhub.io"
        
        # 创建会话并设置请求头
        self.session = requests.Session()
        if self.api_key:
            self.session.headers.update({
                "Authorization": f"Bearer {self.api_key}"
            })

    def upload_image(self, image_path: str) -> Optional[Dict[str, Any]]:
        """
        上传图片到TikHub Sora2服务
        
        Args:
            image_path (str): 图片文件路径
            
        Returns:
            dict: 上传结果，包含media_id等信息
            {
                "id": "media_xxxxxxxxxxxxxxxxxxxxxxxxxx",
                "url": "图片访问链接",
                "kind": "image",
                "width": 图片宽度,
                "height": 图片高度,
                "file_name": "文件名"
            }
        """
        url = f"{self.base_url}/api/v1/sora2/upload_image"
        
        try:
            with open(image_path, 'rb') as image_file:
                files = {'file': image_file}
                response = self.session.post(url, files=files)
                response.raise_for_status()
                return response.json()
        except Exception as e:
            print(f"图片上传失败: {e}")
            return None

    def create_video(self, prompt: str, orientation: str = "portrait", media_id: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        创建视频生成任务
        
        Args:
            prompt (str): 视频描述文本，最多2000字符
            orientation (str, optional): 视频方向，"portrait"(竖屏)或"landscape"(横屏)，默认为"portrait"
            media_id (str, optional): 图片media_id，用于图生视频
            
        Returns:
            dict: 任务创建结果
            {
                "id": "task_xxxxxxxxxxxxxxxxxxxxxxxxxx",
                "priority": 1
            }
        """
        url = f"{self.base_url}/api/v1/sora2/create_video"
        
        data = {
            "prompt": prompt,
            "orientation": orientation
        }
        
        if media_id:
            data["media_id"] = media_id
            
        try:
            response = self.session.post(url, json=data)
            response.raise_for_status()
            result = response.json()
            # 如果返回的数据结构中包含data字段，则直接返回data字段
            # 否则返回整个结果
            print(result)
            return result.get('data', result)
        except Exception as e:
            print(f"创建视频任务失败: {e}")
            return None

    def get_task_status(self, task_id: str) -> Optional[Dict[str, Any]]:
        """
        查询任务状态（可频繁调用）
        
        Args:
            task_id (str): 任务ID，格式如task_01k7dttf0xfx3t7zhhzycjq8e3
            
        Returns:
            dict: 任务状态信息
            {
                "id": "task_id",
                "status": "queued/processing/succeeded/failed",
                "prompt": "视频生成时使用的文本描述",
                "title": "视频标题",
                "progress_pct": 0.5,
                "generations": [
                    {
                        "id": "gen_xxxxxx",
                        "kind": "sora_draft",
                        "url": "视频链接（有水印）",
                        "downloadable_url": "下载链接（有水印）",
                        "width": 视频宽度,
                        "height": 视频高度,
                        "created_at": 时间戳,
                        "prompt": "生成提示词",
                        "encodings": {
                            "source": "源文件",
                            "source_wm": "带水印的源文件",
                            "thumbnail": "缩略图",
                            "md": "中等质量视频",
                            "gif": "预览GIF"
                        }
                    }
                ]
            }
        """
        url = f"{self.base_url}/api/v1/sora2/get_task_status"
        
        params = {
            "task_id": task_id
        }
        
        try:
            response = self.session.get(url, params=params)
            response.raise_for_status()
            result = response.json()
            # 如果返回的数据结构中包含data字段，则直接返回data字段
            # 否则返回整个结果
            return result.get('data', result)
        except Exception as e:
            print(f"查询任务状态失败: {e}")
            return None

    def get_task_detail(self, task_id: Optional[str] = None, generation_id: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        获取任务生成的作品详情（无水印版本）
        
        Args:
            task_id (str, optional): 任务ID，格式如task_01k7e17rnkeh79qnrcdwf5fcfs
            generation_id (str, optional): 生成ID，格式如gen_01k7e1bff9eq6rxe9pntk7xdcf
            
        Returns:
            dict: 作品详细信息（无水印）
            {
                "post": {
                    "id": "作品ID",
                    "text": "作品描述文本",
                    "attachments": [
                        {
                            "url": "无水印视频链接（原始质量）",
                            "downloadable_url": "有水印视频链接",
                            "width": 视频宽度,
                            "height": 视频高度,
                            "encodings": {
                                "thumbnail": "缩略图（无水印）",
                                "md": "中等质量视频（无水印）",
                                "gif": "预览GIF（无水印）"
                            }
                        }
                    ],
                    "like_count": 点赞数,
                    "view_count": 浏览数,
                    "reply_count": 评论数,
                    "remix_count": 混剪数,
                    "posted_at": 发布时间戳,
                    "permalink": "作品永久链接"
                }
            }
        """
        if not task_id and not generation_id:
            raise ValueError("必须提供task_id或generation_id中的一个")
            
        url = f"{self.base_url}/api/v1/sora2/get_task_detail"
        
        params = {}
        if task_id:
            params["task_id"] = task_id
        if generation_id:
            params["generation_id"] = generation_id
            
        try:
            response = self.session.get(url, params=params)
            response.raise_for_status()
            result = response.json()
            # 如果返回的数据结构中包含data字段，则直接返回data字段
            # 否则返回整个结果
            return result.get('data', result)
        except Exception as e:
            print(f"获取任务详情失败: {e}")
            return None

    def wait_for_completion(self, task_id: str, check_interval: int = 2) -> Optional[Dict[str, Any]]:
        """
        等待任务完成并返回最终状态
        
        Args:
            task_id (str): 任务ID
            check_interval (int): 检查间隔（秒），默认为2秒
            
        Returns:
            dict: 最终任务状态
        """
        print(f"等待任务 {task_id} 完成...")
        
        while True:
            status = self.get_task_status(task_id)
            if not status:
                print("获取任务状态失败")
                return None
                
            task_status = status.get("status")
            # 处理progress_pct为None的情况
            progress = status.get("progress_pct") or 0
            progress_percent = progress * 100
            
            print(f"任务状态: {task_status}, 进度: {progress_percent:.1f}%")
            
            if task_status == "succeeded":
                print("任务已完成")
                return status
            elif task_status == "failed":
                print("任务失败")
                return status
            elif task_status not in ["queued", "processing"]:
                print(f"未知任务状态: {task_status}")
                return status
                
            time.sleep(check_interval)